# coding:utf-8
"""
对django的model方法进行扩展
"""

from django.db import connections


def _as_sql(obj, field, connection):
    """
    对原生方法简单实现
    """
    pre_value = field.pre_save(obj, False)
    db_pre_value = field.get_db_prep_save(pre_value, connection)
    return db_pre_value, '%s'


def bulk_update(objs, using='employee', batch_size=None, update_fields=None):
    """
    批量更新
    支持不同model不同字段更新
    没有经过大量测试，慎用
    """
    assert batch_size is None or batch_size > 0

    objs = list(objs)
    if not objs:
        return

    batch_size = batch_size or len(objs)

    connection = connections[using]
    update_sql_template = 'UPDATE {} SET {} WHERE {};'
    for i in range(0, len(objs), batch_size):
        batch_sql_list = []
        batch_params = []
        for item in objs[i:i + batch_size]:
            model_meta = item._meta
            pk_field = model_meta.get_field(model_meta.pk.name)

            # 获取更新字段
            obj_update_fields = getattr(item, '_bulk_update_fields', None) or update_fields
            if not obj_update_fields:
                continue

            # 赋值
            meta_fields = [model_meta.get_field(i) for i in obj_update_fields if not i == pk_field.name]
            values_list = []
            for field in meta_fields:
                value, placeholder = _as_sql(item, field, connection)
                values_list.append('{}={}'.format(field.column, placeholder))
                batch_params.append(value)

            # 主键过滤
            pk_value, pk_placeholder = _as_sql(item, pk_field, connection)
            table_name = model_meta.db_table
            pk_clause = '{}={}'.format(pk_field.column, pk_placeholder)
            batch_params.append(pk_value)
            batch_sql_list.append(update_sql_template.format(table_name, ','.join(values_list), pk_clause))

        if batch_sql_list:
            batch_sql = ' '.join(batch_sql_list)
            with connection.cursor() as cursor:
                cursor.execute(batch_sql, batch_params)


def bulk_update_list(objs, model_meta, using='employee', batch_size=None):
    """
    批量更新, 跳过创建model实例提升速度
    objs: 更新数据, 最终保存到数据库的数据, 字典列表类型
    model_meta: Model._meta
    没有经过大量测试，慎用
    """
    assert batch_size is None or batch_size > 0

    objs = list(objs)
    if not objs or not model_meta:
        return

    batch_size = batch_size or len(objs)
    connection = connections[using]
    update_sql_template = 'UPDATE {} SET {} WHERE {};'
    default_placeholder = '%s'
    table_name = model_meta.db_table
    pk_field = model_meta.pk

    for i in range(0, len(objs), batch_size):
        batch_sql_list = []
        batch_params = []
        for item in objs[i:i + batch_size]:
            # 获取更新字段
            update_fields = item.keys()
            if pk_field.name not in update_fields:
                # 需要根据数据里的主键来更新
                continue

            # 赋值
            meta_fields = [model_meta.get_field(i) for i in update_fields if i != pk_field.name]
            value_list = []
            for field in meta_fields:
                value = field.get_db_prep_save(item.get(field.attname), connection=connection)
                value_list.append('{}={}'.format(field.column, default_placeholder))
                batch_params.append(value)

            # 主键过滤
            pk_value = pk_field.get_db_prep_save(item.get(pk_field.attname), connection=connection)
            pk_clause = '{}={}'.format(pk_field.column, default_placeholder)
            batch_params.append(pk_value)
            batch_sql_list.append(update_sql_template.format(table_name, ','.join(value_list), pk_clause))

        if batch_sql_list:
            batch_sql = ' '.join(batch_sql_list)
            with connection.cursor() as cursor:
                cursor.execute(batch_sql, batch_params)


def bulk_create(objs, model_meta, using='employee', batch_size=None):
    """
    批量创建, 直接生成sql, model自带的批量创建太慢
    objs: 创建数据, 最终保存到数据库的数据, 字典列表类型
    model_meta: Model._meta
    没有经过大量测试，慎用
    """
    assert batch_size is None or batch_size > 0

    objs = list(objs)
    if not objs or not model_meta:
        return

    batch_size = batch_size or len(objs)
    connection = connections[using]
    insert_sql_template = 'INSERT INTO {} ({}) VALUES({});'
    default_placeholder = '%s'
    table_name = model_meta.db_table
    pk_name = model_meta.pk.name

    for i in range(0, len(objs), batch_size):
        batch_sql_list = []
        batch_params = []
        for item in objs[i:i + batch_size]:
            # 获取创建字段
            create_fields = item.keys()
            if pk_name not in create_fields:
                # 数据里必须含有主键
                continue

            # 赋值
            meta_fields = [model_meta.get_field(i) for i in create_fields]
            column_list = []
            value_list = []
            for field in meta_fields:
                value = field.get_db_prep_save(item.get(field.attname), connection=connection)
                column_list.append(field.column)
                value_list.append(default_placeholder)
                batch_params.append(value)

            batch_sql_list.append(insert_sql_template.format(table_name, ','.join(column_list), ','.join(value_list)))

        if batch_sql_list:
            batch_sql = ' '.join(batch_sql_list)
            with connection.cursor() as cursor:
                cursor.execute(batch_sql, batch_params)
